#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""

"""

from pprint import pprint
import copy

from linear_tools import linear_combination
from matrix_tools import multiply, identity, matrix_append_row, matrix_equal, transpose_matrix, resort_mat_row, is_zero_matrix, col_of, matrix_slice_row, show_matrix
from math_tools import float_equal
from dot_cross_product import dot_product


class NoSuitableRow(Exception): pass

class Matrix:
    PIVOT = 1
    FREE  = 0

    # =================== 两个初等变换 =====================
    def _row_switch_trans(self, r1, r2, mat=None):
        ''' 交换行 '''
        if not mat:
            mat = self.matrix

        mat[r1], mat[r2] = mat[r2], mat[r1]

    def _subtract_row(self, eliminating_row_idx, pivot_row_idx, pivot_col_idx, matrix=None):
        ''' k * ri  + 1 * rj '''
        if not matrix:
            matrix = self.matrix

        eliminating_row = matrix[eliminating_row_idx]
        pivot_row       = matrix[pivot_row_idx]

        coefficient = eliminating_row[pivot_col_idx] / pivot_row[pivot_col_idx]
        matrix[eliminating_row_idx] = list( linear_combination(
            [pivot_row, eliminating_row],
            [-coefficient, 1]
        ))
        return coefficient

    def _not_zero(self, r, c):
        return 0 != self.matrix[r][c]

    def _is_zero(self, r, c):
        return 0 == self.matrix[r][c]

    def _record_switch_transform(self, r1, r2):
        print('=================== switch =====================')
        self.transform_list.append('E(%s <--> %s)' % (r1, r2))

        self._row_switch_trans(r1, r2, self.row_elementary_trans)
        self._row_switch_trans(r1, r2, self.row_elementary_switch)

    def _record_elimit_transform(self, pivot_row, cur_row, co):
        '''
            pivot_row: 元来自第几行
            cur_row: 根据pivot_row 消第几行
        '''
        self.transform_list.append('M: (%s * %s) + (1 * %s) ' % (-co, pivot_row, cur_row))

        # 记录初等变换矩阵
        ele_trans = self.row_elementary_trans
        ele_trans[cur_row] = list( linear_combination(
            [ele_trans[pivot_row], ele_trans[cur_row]],
            [-co, 1]
        ))

    def _pivoting(self, pivot_row, pivot_col):
        ''' 如果当前行pivot位置为0, 找到当前行之后的行中pivot_idx/row_idx 不为0 的行, 并和当前行交换.
            如果没有找到, 说明处理的矩阵为singular矩阵: raise NoSuitableRow.
            这个寻找合适的pivot行的过程叫做pivoting
        '''
        if self._not_zero(pivot_row, pivot_col):
            return

        for cand_row in range(pivot_row+1, self.rows):
            if self._not_zero(cand_row, pivot_col):
                self._row_switch_trans(pivot_row, cand_row)

                # 记录交换行操作
                self._record_switch_transform(pivot_row, cand_row)
                return

        raise NoSuitableRow()

    def _eliminate_row(self, eliminating_row, pivot_row, pivot_col):
        if self._is_zero(eliminating_row, pivot_col):
            return

        coefficient = self._subtract_row(eliminating_row, pivot_row, pivot_col)

        # 记录消元的初等变换过程
        self._record_elimit_transform(pivot_row, eliminating_row, coefficient)

    def _normalization(self, row, col):
        ''' 对 row做归一化处理, 基于(row, col)来做

            示例输入: 2 2 3
                      0 2 5
                      0 0 6

            示例输出: 2 2 3 ------> 1, 1, 1.5
                      0 2 5 ------> 0, 1, 2.5
                      ...
        '''
        co = self.matrix[row][col]
        assert co
        if co == 1:
            return

        self.matrix[row]               = [value/co for value in self.matrix[row]]
        self.row_elementary_trans[row] = [value/co for value in self.row_elementary_trans[row]]
        # self.row_elementary_trans[row][row] = 1/co

    def _eliminate_pivot(self, row, col):
        ''' 处理第几个pivot '''
        self._pivoting(row, col)    # 保证pivot_idx维护行这个位置不为0
        self._normalization(row, col)

        for eliminating_row in range(0, self.rows):
            '''  NOTE: 这里大胆的把当前行之前的行和当前行之后的行都消去, 省掉了最后一步back-substitution '''
            if eliminating_row == row:
                # 忽略pivot row
                continue
            self._eliminate_row(eliminating_row, row, col)

    def _init_matrix(self, matrix, inplace=True):
        ''' 对aug_matrix 增光矩阵进行高斯消元, 可以用来求逆矩阵, 也可以解方程组 '''
        self.matrix = matrix
        if not inplace:
            self.matrix = copy.deepcopy(matrix)

        self.rows = len(matrix)
        self.cols = len(matrix[0])

        self.transform_list = []                         # 用"字符串"描述消元过程
        self.row_elementary_switch = identity(self.rows) # 初等交换矩阵: 每次初等变换, 记录交换两行的操作, 记为S
        self.row_elementary_trans = identity(self.rows)  # 初等变换矩阵, 随着每次初等变换操作, 最后变为EA = U 的"E"(不包括交换操作)

    def _identify_pivot_frees(self, reduced):
        ''' 寻找reduced矩阵中的pivot和free列 '''
        pivot_free_indicator = [self.FREE] * self.cols       # 默認填充FREE

        row = 0
        for col in range(self.cols):
            if float_equal(reduced[row][col], 1):
                pivot_free_indicator[col] = self.PIVOT
                row += 1

            if row >= self.rows:
                break

        return pivot_free_indicator

    def _build_null_basis(self, reduced, free_idxs, pivot_idxs):
        ''' 已知reduced矩阵, free_idxs, pivot_idxs, 求null_space 的基 '''
        if len(free_idxs) == 0:
            null_space_basis = [[0]] * self.cols
            return null_space_basis

        free_n  = len(free_idxs)
        pivot_n = self.cols - free_n

        # 1: free_part in pivot_n row
        F_part = []
        for row in self.matrix[0:pivot_n]:
            free_part_row = [-row[free_col] for free_col in free_idxs]
            F_part.append(free_part_row)
        # print('F_part %s' % F_part)

        # 2: identity part in free_n row
        I_part = identity(free_n)

        # 3: F + I ====> null-space-basic
        null_space_basis = matrix_append_row(F_part, I_part, inplace=True)
        # basis是由列組成的
        # null_space_basis = transpose_matrix(null_space_basis)

        # 这里比较费解一些: 原来的索引, 在现在的哪一行. 记录下来, 用于重新换null_space_basis行归位.
        resort_idxs = [0] * self.cols
        new_null_matrix_idxs = pivot_idxs + free_idxs
        for now_idx, old_idx in enumerate(new_null_matrix_idxs):
            resort_idxs[old_idx] = now_idx

        # show_matrix(null_space_basis, 'null')
        # print("p idxs %s" % (pivot_idxs ))
        # print("f idxs %s" % (free_idxs ))
        # print("p+f idxs %s" % (pivot_idxs + free_idxs))
        null_space_basis = resort_mat_row(null_space_basis, resort_idxs)

        return null_space_basis

    # =================== outside interface =====================
    def rref(self, mat, inplace=False):
        ''' 对matrix矩阵进行高斯消元, 获取row-reduce-echelon-form形式.
            可以用来求逆矩阵, 也可以解方程组
            允许应用在后面扩展
        '''
        self._init_matrix(mat, inplace)

        row = 0
        for col in range(self.cols):
            try:
                self._eliminate_pivot(row, col)
            except NoSuitableRow as e:
                pass
            else:
                row += 1

            if row >= self.rows:
                break

        return self.row_elementary_trans, self.matrix

    def four_space_basis(self, mat):
        ''' 求取 col space basis
                null space basis
        '''
        ele_trans, reduced = self.rref(mat)
        show_matrix(reduced, 'reduced in four_space_basis')
        show_matrix(ele_trans, 'ele_trans in four_space_basis')

        ele_trans_A = multiply(ele_trans, mat)
        show_matrix(ele_trans_A, 'ele_trans * mat == reduced')
        matrix_equal(ele_trans_A, reduced)

        pivot_free_indicator = self._identify_pivot_frees(reduced)
        # print('pivot free indicator %s' % pivot_free_indicator)

        # ===================  null space basis(是竖着的, 没有转置) =====================
        # if all(pivot_free_indicator):
        #     # 列全部线性无关
        #     null_space_basis = [[0]] * self.cols
        #     # return
        # else:
        pivot_idxs = []
        free_idxs  = []
        for col_idx, indicator in enumerate(pivot_free_indicator):
            if indicator:
                pivot_idxs.append(col_idx)
            else:
                free_idxs.append(col_idx)

        null_space_basis = self._build_null_basis(reduced, free_idxs, pivot_idxs)

        A_nullbasis = multiply(mat, null_space_basis)
        # A_nullbasis = multiply(reduced, null_space_basis)
        assert is_zero_matrix(A_nullbasis)
        show_matrix(null_space_basis, 'null basis')
        show_matrix(A_nullbasis, 'A * null_space_basis = 0 matrix')

        print('=======================================  ===========================================')
        # ===================  column space basis =====================
        column_space_basis = []
        for col in pivot_idxs:
            column_space_basis.append(col_of(mat, col))
        column_space_basis = transpose_matrix(column_space_basis)
        show_matrix(column_space_basis, "column_space_basis")

        # =================== row basis =====================
        rank = len(pivot_idxs)
        # row_space_basis = matrix_slice_row(reduced, range(rank))
        row_space_basis = reduced[0:rank]
        show_matrix(row_space_basis, 'row space basis')

        # =================== left null basis =====================
        if rank == self.rows:
            left_null_basis = [[0]] * self.rows
        else:
            left_null_basis = ele_trans[rank:]
        show_matrix(left_null_basis, 'left_null space basis')
        left_A = multiply(left_null_basis, mat)
        show_matrix(left_A, 'left_null_basis * A = zero_vector')

        # =================== orthogonal complement  =====================
        null_row = multiply(row_space_basis, null_space_basis)
        show_matrix(null_row, 'null_row')
        left_null_column = multiply(left_null_basis, column_space_basis)
        show_matrix(left_null_column, 'left_null_column')

        # dot_product
        null_space_basis = transpose_matrix(null_space_basis)
        column_space_basis = transpose_matrix(column_space_basis)

        for n in null_space_basis:
            for r in row_space_basis:
                print(int (dot_product(n, r)))

        print('===================  =====================')
        for n in left_null_basis:
            for r in column_space_basis:
                print(int (dot_product(n, r)))

        print()

        d = {
            'col' : column_space_basis,
            'row' : row_space_basis,
            'null' : null_space_basis,
            'left' : left_null_basis,
        }
        return d


# def show_matrix(mat, msg=''):
#     print('=================== mat: %s =====================' % msg)
#     for r in mat:
#         print(r)
#     print()

def test_reduced_form():
    # A = [
    #     [1, 2, 3, 4, 5],
    #     [1, 2, 3, 4, 5],
    #     [55, 13, 323, 18, 901],
    #     [13, 8, 3, 9, 18],
    #     [3234, 200, 3, 9, 18],
    #     [135, 215, 3, 9, 18],
    #     [15123, 20, 3, 9, 18],
    # ]
    A = [
        [1, 2, 3, 4, 5],
        [0, 0, 3, 6, 8],
    ]

    # ===================  =====================
    m = Matrix()
    row_elementary_trans, reduced = m.rref(A)

    show_matrix(A, 'A')
    show_matrix(reduced, 'reduced')
    show_matrix(row_elementary_trans, 'row_elementary_trans matrix')

    trans_a = multiply(row_elementary_trans, A)
    show_matrix(trans_a, 'trans_mat * A == reduced')
    matrix_equal(reduced, trans_a)      # check equal

def test_null_space():
    # A = [
    #     [1, 2, 3, 4],
    #     [3, 2, 3, 4],
    #     [0, 9, 4, 10],
    #     [5, 0, 4, 10],
    # ]

    # A = [
    #     [1, 2, 3, 3, 2, 5],
    #     [3, 2, 3, 3, 2, 5],
    #     [0, 9, 4, 4, 9, 13],
    #     [5, 0, 4, 4, 0, 5],
    # ]

    A = [
        [1, 2, 3, 3],
        [0, 9, 4, 5],
        [0, 9, 4, 5],
        [3, 2, 3, 5],
        [3, 11, 7, 10],
        [3, 2, 3, 5],
    ]

    # A = [
    #     [1, 0, 0],
    #     [0, 1, 0],
    #     [0, 0, 1],
    #     [0, 0, 0],
    #     [0, 0, 0],

    #     [2, 0, 0],
    #     [0, 2, 0],
    #     [0, 0, 2],
    #     [0, 0, 0],
    #     [0, 0, 0],
    # ]

    m = Matrix()
    m.four_space_basis(A)


def test():
    # test_reduced_form()
    test_null_space()

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
